import os
import shutil

import numpy as np
import cv2

# import pandas as pd

labels = {'Anomalies': 0, 'Cyst': 1, 'Inflammation': 2, 'Tumor': 3, 'Vascular': 4}


def load_numpy_from_subdirs(base_path):
    data = {}
    for subdir in os.listdir(base_path):
        subdir_path = os.path.join(base_path, subdir)
        if os.path.isdir(subdir_path):
            numpy_files = [f for f in os.listdir(subdir_path) if f.endswith('.npy')]
            category_data = []
            for numpy_file in numpy_files:
                numpy_file_path = os.path.join(subdir_path, numpy_file)
                data_array = np.load(numpy_file_path)
                category_data.append(data_array)
            data[subdir] = category_data
    return data


def read_data(base_path='data/train'):
    data = load_numpy_from_subdirs(base_path)

    # 查看数据结构
    for category, arrays in data.items():
        print(f"Category: {category}, Number of Arrays: {len(arrays)}")
    print(data['Anomalies'][0].shape)

    return data


def predict(model, image):
    image = np.expand_dims(image, axis=0)  # 扩展维度以匹配模型输入
    prediction = model.predict(image)
    print(prediction)
    predicted_class = np.argmax(prediction, axis=1)  # 获取预测类别的索引
    return predicted_class


def model_predict(test_path='test', model_path='cnn_model.h5'):
    import tensorflow as tf
    # 加载模型
    loaded_model = tf.keras.models.load_model(model_path)

    # 假设 test_image 是你想要预测的图片
    pres_list = []

    for test_img_name in os.listdir(test_path):
        test_img_path = test_path + '/' + test_img_name
        predicted_class_index = predict(loaded_model, np.load(test_img_path))
        predicted_class_name = list(labels.keys())[list(labels.values()).index(predicted_class_index)]
        print(f"Predicted Class: {predicted_class_name}")
        pres_list.append(predicted_class_name)
    pres_list = []
    uuid_list = []

    for idx, test_img_name in enumerate(os.listdir(test_path), start=1):
        test_img_path = os.path.join(test_path, test_img_name)
        predicted_class_index = predict(loaded_model, np.load(test_img_path))
        predicted_class_name = list(labels.keys())[list(labels.values()).index(predicted_class_index)]
        uuid = test_img_name.replace('.npy', '')
        print(f"Predicted Class: {test_img_name, predicted_class_name, uuid}")
        uuid_list.append(uuid)  # 使用 enumerate 函数自动递增
        pres_list.append(predicted_class_name)

    # 创建 DataFrame
    df = pd.DataFrame({
        'uuid': uuid_list,
        'label': pres_list
    })
    # 保存 DataFrame 到 CSV 文件
    df['uuid'] = df['uuid'].astype(int)
    df = df.sort_values(by='uuid', ascending=True)
    df.to_csv('{}_predictions.csv'.format(model_path.replace('.h5', '')), index=False)


def resize_and_save_npy(root_dir, output_dir, target_size=(224, 224)):
    import cv2
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 遍历根目录
    for dirpath, _, filenames in os.walk(root_dir):
        for filename in filenames:
            if filename.endswith('.npy'):
                file_path = os.path.join(dirpath, filename)
                # 读取.npy文件
                image = np.load(file_path)

                # 检查图像通道数，确保它是3通道，如果不是，可以进行相应处理
                if image.shape[-1] != 3:
                    continue  # 或者你可以添加转换代码来适配这个条件

                # 调整大小
                resized_image = cv2.resize(image, target_size)

                # 构建输出路径
                rel_path = os.path.relpath(dirpath, root_dir)
                output_path = os.path.join(output_dir, rel_path)

                if not os.path.exists(output_path):
                    os.makedirs(output_path)

                output_file_path = os.path.join(output_path, filename)

                # 保存调整大小后的图像
                np.save(output_file_path, resized_image)
                print(f"Saved resized image to {output_file_path}")


def convert_png_to_npy_with_opencv(directory, output_directory, target_size=(224, 224)):
    import cv2
    if not os.path.exists(output_directory):
        os.makedirs(output_directory)

    for root, dirs, files in os.walk(directory):
        for file in files:
            if file.lower().endswith('.png'):
                file_path = os.path.join(root, file)
                # 使用OpenCV读取图片
                image = cv2.imread(file_path)
                image = cv2.resize(image, target_size)
                # OpenCV默认使用BGR颜色空间，如果需要RGB可以转换
                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
                # 构建输出路径
                relative_path = os.path.relpath(file_path, directory)
                base_name = os.path.splitext(relative_path)[0]
                output_path = os.path.join(output_directory, base_name + '.npy')
                # 确保输出目录存在
                output_subdir = os.path.dirname(output_path)
                if not os.path.exists(output_subdir):
                    os.makedirs(output_subdir)
                # 保存.npy文件
                np.save(output_path, image)
                print(f"Saved: {output_path}")


def convert_npy_to_png(directory):
    # 遍历目录及其子目录
    for root, dirs, files in os.walk(directory):
        for file in files:
            if file.endswith(".npy"):
                npy_path = os.path.join(root, file)
                png_path = os.path.splitext(npy_path)[0] + ".png"

                # 加载npy文件
                img_array = np.load(npy_path)

                # 如果npy数据是浮点型，将其归一化为[0, 255]区间并转换为uint8
                if img_array.dtype == np.float32 or img_array.dtype == np.float64:
                    img_array = (img_array - np.min(img_array)) / (np.max(img_array) - np.min(img_array)) * 255
                    img_array = img_array.astype(np.uint8)

                # 保存为png文件
                cv2.imwrite(png_path, img_array)
                print(f"Converted: {npy_path} -> {png_path}")


def split_data(source_dir, base_test_dir, test_ratio):
    # 遍历源目录下的所有子目录
    for subdir in os.listdir(source_dir):
        subdir_path = os.path.join(source_dir, subdir)

        # 确保子目录是一个目录
        if os.path.isdir(subdir_path):
            # 创建对应的测试集目录
            dest_dir = os.path.join(base_test_dir, subdir)
            if not os.path.exists(dest_dir):
                os.makedirs(dest_dir)

            # 获取子目录中所有png文件
            files = [f for f in os.listdir(subdir_path) if f.endswith('.png')]

            # 计算要移动到测试集的文件数量
            n_test = int(len(files) * test_ratio)

            # 随机选择文件
            test_files = np.random.choice(files, n_test, replace=False)

            # 移动选中的文件到新目录
            for file in test_files:
                src_path = os.path.join(subdir_path, file)
                dest_path = os.path.join(dest_dir, file)
                shutil.move(src_path, dest_path)
                print(f'Moved: {src_path} -> {dest_path}')


if __name__ == '__main__':
    # root_dir = 'train'
    # output_dir = 'train_224'
    # resize_and_save_npy(root_dir, output_dir)
    # root_dir = 'test'
    # output_dir = 'test_224'
    # resize_and_save_npy(root_dir, output_dir)
    # convert_png_to_npy_with_opencv(directory='E:/data/csb_img/merge_data/train',
    #                                output_directory='E:/data/csb_img/exc_data')
    source_directory = 'data/train'
    test_directory = 'data/test'

    # 设置测试集占比
    test_ratio = 0.10

    # 执行数据分割
    split_data(source_directory, test_directory, test_ratio)
